# Port scaffold of mod_save.f90

module OutputWriter

using Printf
using NCDatasets

"""
    save_region_products(outdir::AbstractString, domain)

Writes nbox_xy.txt, lsm_box.txt, area_box.txt, and hloc_box.txt similar to Fortran outputs.
"""
function save_region_products(outdir::AbstractString, domain)
    isdir(outdir) || mkpath(outdir)
    # nbox_xy.txt
    open(joinpath(outdir, "nbox_xy.txt"), "w") do io
        for i in 1:domain.nxregrid
            @printf(io, "%s\n", join(domain.nbox_xy[i, :], ' '))
        end
    end
    # lsm_box.txt
    open(joinpath(outdir, "lsm_box.txt"), "w") do io
        for i in 1:domain.nxregrid
            @printf(io, "%s\n", join(domain.lsm_box[i, :], ' '))
        end
    end
    # area_box.txt
    open(joinpath(outdir, "area_box.txt"), "w") do io
        for n in 1:length(domain.area_box)
            @printf(io, "%11.4E\n", domain.area_box[n])
        end
    end
    # hloc_box.txt
    open(joinpath(outdir, "hloc_box.txt"), "w") do io
        for n in 1:length(domain.hloc)
            @printf(io, "%5.1f\n", domain.hloc[n])
        end
    end
    return nothing
end

"""
    save_obs(path::String, obs::Observations, y_mod::Vector{Float64}=Float64[])

Save observation data and model-observation comparison.
"""
function save_obs(path::String, obs, y_mod::Vector{Float64}=Float64[])
    isdir(dirname(path)) || mkpath(dirname(path))

    open(path, "w") do io
        # Header
        println(io, "# Station JulianDate DateTime Observed ModeledObs Error Residual")

        for i in 1:length(obs.concentrations)
            station = length(obs.stations) >= i ? obs.stations[i] : "UNK"
            juldate = length(obs.jdates) >= i ? obs.jdates[i] : 0.0
            datetime_str = length(obs.datetimes) >= i ? string(obs.datetimes[i]) : "missing"
            observed = obs.concentrations[i]
            modeled = length(y_mod) >= i ? y_mod[i] : 0.0
            error = length(obs.measurement_errors) >= i ? obs.measurement_errors[i] : 0.0
            residual = modeled - observed

            @printf(io, "%s %12.5f %s %10.4f %10.4f %8.4f %10.4f\n",
                   station, juldate, datetime_str, observed, modeled, error, residual)
        end
    end

    @info "Observations saved to $path"
end

"""
    save_state(path::String, state::InversionState, domain::Domain, covariance::CovarianceMatrix)

Save state vector and uncertainty information.
"""
function save_state(path::String, state, domain, covariance)
    isdir(dirname(path)) || mkpath(dirname(path))

    # Get physical state
    chi_state = get_physical_state(state)

    open(path, "w") do io
        # Header
        println(io, "# Box Longitude Latitude State Prior Increment Uncertainty")

        for i in 1:domain.nbox
            # Find grid coordinates for this box
            ix, jy = find_box_indices(domain, i)
            lon = domain.rllx + (ix - 1) * domain.rdx
            lat = domain.rlly + (jy - 1) * domain.rdy

            current_state = length(chi_state) >= i ? chi_state[i] : 0.0
            prior_state = length(state.chi_prior) >= i ? state.chi_prior[i] : 0.0
            increment = current_state - prior_state
            uncertainty = 0.0  # Would compute from posterior covariance if available

            @printf(io, "%6d %10.4f %10.4f %12.6E %12.6E %12.6E %12.6E\n",
                   i, lon, lat, current_state, prior_state, increment, uncertainty)
        end
    end

    @info "State vector saved to $path"
end

"""
    save_analysis_netcdf(filepath::String, state::InversionState, domain::Domain,
                        covariance::CovarianceMatrix, obs::Observations, result::InversionResult)

Save analysis results to NetCDF format compatible with Fortran output.
"""
function save_analysis_netcdf(filepath::String, state, domain, covariance, obs, result)
    isdir(dirname(filepath)) || mkpath(dirname(filepath))

    Dataset(filepath, "c") do ds
        # Define dimensions
        defDim(ds, "longitude", domain.nxregrid)
        defDim(ds, "latitude", domain.nyregrid)
        defDim(ds, "time", 1)  # Single time step for now
        defDim(ds, "nbox", domain.nbox)
        defDim(ds, "modes", covariance.n_modes)

        # Define coordinate variables
        lon_var = defVar(ds, "longitude", Float64, ("longitude",))
        lat_var = defVar(ds, "latitude", Float64, ("latitude",))
        time_var = defVar(ds, "time", Float64, ("time",))

        # Define data variables
        state_var = defVar(ds, "state_vector", Float64, ("nbox",))
        prior_var = defVar(ds, "prior_state", Float64, ("nbox",))
        increment_var = defVar(ds, "state_increment", Float64, ("nbox",))

        control_var = defVar(ds, "control_vector", Float64, ("modes",))
        eigenvals_var = defVar(ds, "eigenvalues", Float64, ("modes",))

        # Fill coordinate variables
        lon_var[:] = [domain.rllx + (i-1)*domain.rdx for i in 1:domain.nxregrid]
        lat_var[:] = [domain.rlly + (j-1)*domain.rdy for j in 1:domain.nyregrid]
        time_var[:] = [1.0]  # Placeholder

        # Fill data variables
        chi_state = get_physical_state(state)
        state_var[:] = chi_state[1:domain.nbox]
        prior_var[:] = state.chi_prior[1:domain.nbox]
        increment_var[:] = chi_state[1:domain.nbox] .- state.chi_prior[1:domain.nbox]

        control_var[:] = state.phi
        eigenvals_var[:] = covariance.eigenvalues

        # Add attributes
        ds.attrib["title"] = "FLEXINVERT Analysis Results"
        ds.attrib["method"] = result.method
        ds.attrib["final_cost"] = result.final_cost
        ds.attrib["iterations"] = result.iterations
        ds.attrib["n_observations"] = length(obs.concentrations)
        ds.attrib["n_modes"] = covariance.n_modes
        ds.attrib["explained_variance"] = covariance.explained_variance

        lon_var.attrib["units"] = "degrees_east"
        lat_var.attrib["units"] = "degrees_north"
        state_var.attrib["units"] = "flux_units"
        state_var.attrib["long_name"] = "Analysis state vector"
    end

    @info "Analysis NetCDF saved to $filepath"
end

"""
    save_inversion_summary(path::String, result::InversionResult, obs::Observations,
                          covariance::CovarianceMatrix)

Save a text summary of the inversion results.
"""
function save_inversion_summary(path::String, result, obs, covariance)
    isdir(dirname(path)) || mkpath(dirname(path))

    open(path, "w") do io
        println(io, "FLEXINVERT Inversion Summary")
        println(io, "=" ^ 50)
        println(io)

        println(io, "Method: $(result.method)")
        println(io, "Success: $(result.success)")
        println(io, "Convergence: $(result.convergence_info)")
        println(io, "Iterations: $(result.iterations)")
        println(io, "Final cost: $(result.final_cost)")
        println(io)

        println(io, "Observations: $(length(obs.concentrations))")
        if length(obs.concentrations) > 0
            println(io, "Concentration range: $(minimum(obs.concentrations)) to $(maximum(obs.concentrations))")
            println(io, "Error range: $(minimum(obs.measurement_errors)) to $(maximum(obs.measurement_errors))")
        end
        println(io)

        println(io, "Prior covariance:")
        println(io, "  Retained modes: $(covariance.n_modes)")
        println(io, "  Explained variance: $(round(covariance.explained_variance*100, digits=1))%")
        println(io, "  Total variance: $(covariance.total_variance)")
        println(io)

        if !isempty(result.cost_history)
            println(io, "Cost function history:")
            for (i, cost) in enumerate(result.cost_history)
                println(io, "  Iter $i: $cost")
            end
        end
    end

    @info "Inversion summary saved to $path"
end

# Helper functions

"""
    find_box_indices(domain::Domain, box_id::Int) -> Tuple{Int, Int}

Find grid indices (ix, jy) for a given box ID.
"""
function find_box_indices(domain, box_id::Int)
    for jy in 1:domain.nyregrid
        for ix in 1:domain.nxregrid
            if domain.nbox_xy[ix, jy] == box_id
                return ix, jy
            end
        end
    end
    return 1, 1  # Fallback
end

"""
    get_physical_state(state)

Extract physical state from InversionState object.
"""
function get_physical_state(state)
    if hasfield(typeof(state), :chi)
        return state.chi
    else
        return Float64[]
    end
end

end # module
